# 这里是条件GAN的代码。# 这里要多传入一个labels,这个样子就可以输出你想要的数字。
import time
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import utils, datasets, transforms
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML
import torchvision

device = 'cuda' if torch.cuda.is_available() else 'cpu'


class generator(nn.Module):
    def __init__(self):
        super(generator, self).__init__()
        self.noise_net = nn.Sequential(
            nn.ConvTranspose2d(100, 256, 4, 1, 0, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)  # 这里是对噪声进行处理
        )
        self.label_net = nn.Sequential(
            nn.ConvTranspose2d(10, 256, 4, 1, 0, bias=False),
            nn.BatchNorm2d(256),  # 这里是对标签进行处理。
            nn.ReLU(inplace=True)
        )
        self.model1 = nn.Sequential(
            nn.ConvTranspose2d(256 + 256, 256, 4, 2, 1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        self.model2 = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )
        self.model3 = nn.Sequential(
            nn.ConvTranspose2d(128, 1, 4, 2, 1, bias=False),
            nn.Tanh()  # 在这里就不要使用inplace = False.
        )

    def forward(self, noise, labels):
        noise = self.noise_net(noise)
        labels = self.label_net(labels)
        print(labels)
        sum_net = torch.cat((noise, labels), dim=1)  # 这里将两个网络合成到一起。
        sum_net = self.model1(sum_net)
        sum_net = self.model2(sum_net)
        sum_net = self.model3(sum_net)
        return sum_net


class discriminator(nn.Module):
    def __init__(self):
        super(discriminator, self).__init__()
        self.images = nn.Sequential(
            nn.Conv2d(1, 64, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.labels = nn.Sequential(
            nn.Conv2d(10, 64, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.model1 = nn.Sequential(
            nn.Conv2d(64 + 64, 256, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.model2 = nn.Sequential(
            nn.Conv2d(256, 512, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        )
        self.model3 = nn.Sequential(
            nn.Conv2d(512, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()  # 最后使用sigmoid函数作为最终的输出。
        )

    def forward(self, images, labels):
        images = self.images(images)
        labels = self.labels(labels)
        # print(labels)
        sum_mat = torch.cat((images, labels), dim=1)  # 这里的dim = 1意思就是按照第1维进行拼接，dim = 0是batchsize dim  =1是channel数量。
        sum_mat = self.model1(sum_mat)
        sum_mat = self.model2(sum_mat)
        sum_mat = self.model3(sum_mat)
        return sum_mat


def hyperparameters():  # 在这里是超参数的设置
    learning_rate = 0.0001
    epoch = 100
    batch_size = 256
    nc = 1
    device = 'cuda' if torch.cuda.is_available() else 'cpu'  # 弄到GPU上面。
    num_work = 4  # 这里是在使用dataloader的时候可以设置一下num_work的数量，主要取决于cpu的性能好坏。

    return learning_rate, epoch, batch_size, nc, device, num_work


def get_data(batch_size):  # 这里是获取数据集的函数,这里要传入的参数是batch_size.
    trans = transforms.Compose([
        transforms.Resize(32),  # 将数据扩大。
        transforms.ToTensor(),
        transforms.Normalize(0.5, 0.5)
    ])
    train_data = torchvision.datasets.MNIST('data',
                                            train=True,
                                            transform=trans,
                                            download=True)
    test_data = torchvision.datasets.MNIST('data',
                                           train=False,
                                           transform=trans,
                                           download=True)
    data = train_data + test_data  # 这里由于是GAN网络，没有训练集以及测试集，所以两个数据集全部可以拿来使用
    dataloader = DataLoader(dataset=data, batch_size=batch_size, shuffle=True,
                            drop_last=True)  # 这里的drop_last表示的是将后面的几个数据提取出来。
    return dataloader


def weight_init(m):  # 这里是对网络进行权重的初始化，这个要学习一下每一步的操作，我觉得对未来构建网络有很大的帮助。
    classname = m.__class__.__name__  # 这里要解释一下这个代码的意思，首先m表示的是一个网络
    # print(classname)  # 这个clasname表示的是你每一次使用了哪些东西，比如Linear，Conv之类的
    if classname.find('Conv') != -1:  # 这里使用了find这个方法，就是找出你使用了哪些东西，只要这些权重不等于-1，则就改变权重，也就是权重初始化。
        nn.init.normal_(m.weight.data, mean=0.0, std=0.02)  # 这里是将其满足均值为0，标准差为0.2的正态分布的随机初始化init就是初始化的意思。
    elif classname.find('BatchNorm') != -1:  # 这里是对BatchNorm进行初始化参数。
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)  # 如果要对这个进行运用的话，net.apply(wight_init) 这里和R语言十分相似。apply就是使用一个函数。


def show_one_image(noise, label):
    noise = noise.to(device)
    # label = label.to(device)  # 这里是将数据转化到GPU上面。
    label_1hots = torch.zeros(10, 10)
    for i in range(10):
        label_1hots[i, i] = 1
    label_1hots = label_1hots.view(10, 10, 1, 1)  # 这里将labels转化成了一个10通道的向量，前面那个10表示的应该是分类的数量
    # print(label_1hots[1].shape)  # 在这里创造一个10维的one_hot编码，到了训练的时候就可以直接锁定这个对应的编码。
    # print(label_1hots.shape)
    label = label_1hots[label]
    label = label.to(device)
    G = torch.load('./G_CGAN_new.pth')
    G = G.to(device)
    image = G(noise, label)  # 这里就是生成的图像。
    image = torch.squeeze(image)
    Toimage = transforms.ToPILImage()
    image = Toimage(image)  # 这里转化为PIL的格式。
    print('Please review the resulting image!')
    image.show()
